package dssh

import (
	"io"
	"log"
	"net"
	"time"

	"golang.org/x/crypto/ssh"
)

/*本地转发*/
// 转发
func sForward(serverAddr string, remoteAddr string, localConn net.Conn, config *ssh.ClientConfig) {
	// 设置sshClientConn
	sshClientConn, err := ssh.Dial("tcp", serverAddr, config)
	if err != nil {
		log.Printf("ssh.Dial failed: %s", err)
	}

	// 设置Connection
	sshConn, err := sshClientConn.Dial("tcp", remoteAddr)

	// 将localConn.Reader复制到sshConn.Writer
	go func() {
		_, err = io.Copy(sshConn, localConn)
		if err != nil {
			log.Printf("io.Copy failed: %v", err)
		}
	}()

	// 将sshConn.Reader复制到localConn.Writer
	go func() {
		_, err = io.Copy(localConn, sshConn)
		if err != nil {
			log.Printf("io.Copy failed: %v", err)
		}
	}()
}

/*
username: 用户名
password: 密码
serverAddr: 服务器地址
remoteAddr: 远程地址
localAddr: 本地地址

建立一条从本地到服务器的隧道
*/
func Tunnel(username string, password string, serverAddr string, remoteAddr string, localAddr string) {
	// 设置SSH配置
	log.Printf("建立一条从本地 %s 到（%s）服务器的 %s 的隧道\n", localAddr, serverAddr, remoteAddr)
	config := &ssh.ClientConfig{
		User: username,
		Auth: []ssh.AuthMethod{
			ssh.Password(password),
		},
		Timeout: 30 * time.Second,
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			return nil
		},
	}

	// 设置本地监听器
	localListener, err := net.Listen("tcp", localAddr)
	if err != nil {
		log.Printf("net.Listen failed: %v\n", err)
	}

	for {
		// 设置本地
		localConn, err := localListener.Accept()
		if err != nil {
			log.Printf("localListener.Accept failed: %v\n", err)
		}
		go sForward(serverAddr, remoteAddr, localConn, config)
	}
}

/*远程转发*/
/*
username: 用户名
password: 密码
serverAddr: 服务器地址
remoteAddr: 远程地址
localAddr: 本地地址

建立一条从服务器到本地的远程转发
*/
func TunnelRemote(username string, password string, serverAddr string, remoteAddr string, localAddr string) {
	// 设置SSH配置
	log.Printf("建立一条从（%s）服务器里的 %s 到 %s 的远程转发\n", serverAddr, remoteAddr, localAddr)
	config := &ssh.ClientConfig{
		User: username,
		Auth: []ssh.AuthMethod{
			ssh.Password(password),
		},
		Timeout: 30 * time.Second,
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			return nil
		},
	}

	// 设置SSH客户端
	sshClientConn, err := ssh.Dial("tcp", serverAddr, config)
	if err != nil {
		log.Printf("ssh.Dial failed: %s", err)
	}

	// 设置远程服务器上remoteAddr的监听器
	remoteListener, err := sshClientConn.Listen("tcp", remoteAddr)
	if err != nil {
		log.Printf("sshClientConn.Listen failed: %v\n", err)
	}
	log.Println("remoteListener: ", remoteListener.Addr())

	for {
		// 设置远程服务器上的remoteAddr的连接
		remoteConn, err := remoteListener.Accept()
		if err != nil {
			log.Printf("remoteListener.Accept failed: %v\n", err)
		}
		go sForwardRemote(remoteConn, localAddr)
	}
}

// 转发，将新的连接转发到本地
func sForwardRemote(remoteConn net.Conn, localAddr string) {
	// 设置本地的连接
	localConn, err := net.Dial("tcp", localAddr)
	if err != nil {
		log.Printf("net.Dial failed: %v\n", err)
	}
	log.Println("Connected to ", localAddr)

	// 将remoteConn.Reader复制到localConn.Writer
	go func() {
		_, err = io.Copy(localConn, remoteConn)
		if err != nil {
			log.Printf("io.Copy failed: %v", err)
		}
	}()

	// 将localConn.Reader复制到remoteConn.Writer
	go func() {
		_, err = io.Copy(remoteConn, localConn)
		if err != nil {
			log.Printf("io.Copy failed: %v", err)
		}
	}()
}
